# -*- coding:utf-8 -*-
import re
import numpy as np
from metlib.datetime import T, TD
from metlib.kits import *
from weblib.common.file_transporter import *
from filecache.filecache_manager import *
from matplotlib.cm import *

merra2_py_summary_sub_varnames = ['wspd50m', 'wpd50m', 't2m', 'qv2m', 'psfc', 'slp']
merra2_pt_varnames = ['wspd', 'wdir', 'wpd', 't', 'qv', 'psfc', 'slp']
merra2_py_varnames = ['summary', 'rose', 'wpdrose', 'dist']
merra2_ry_varnames = ['wspd', 'wspd_max', 'wpd', 't', 't_max', 't_min', 'qv', 'psfc', 'slp']
merra2_pm_varnames = ['wspd', 'wpd', 't', 'qv', 'psfc', 'slp']
merra2_pd_varnames = ['wspd', 'wpd', 't', 'qv', 'psfc', 'slp']

merra2_years = [str(y) for y in range(1979, 2017)]

merra2_dx = 0.625
merra2_dy = 0.5

merra2_full_lons = np.arange(-180.0, 179.34, merra2_dx)
merra2_full_lats = np.arange(-90.0, 90.1, merra2_dy)

merra2_zh_dataset_d = {
    'merra2': 'Merra2'
}

merra2_zh_subset_d = {
   "PT": u"时间序列",
   "PY": u"年统计值",
   "PM": u"月统计值",
   "PD": u"日变化",
   "RY": u"年统计值",
}

merra2_zh_varname_d = {
    "wspd": u"风速",
    "wspd_max": u"最大风速",
    "wdir": u"风向",
    "wpd": u"风功率密度",
    "t": u"温度",
    "t_max": u"最高温度",
    "t_min": u"最低温度",
    "psfc": u"地面气压",
    "slp": u"海平面气压",
    "qv": u"比湿",
    "summary": u"汇总统计",
    "rose": u"风玫瑰",
    "wpdrose": u"风功率玫瑰",
    "dist": u"风速分布",
}

# 单位的字典
merra2_units_d = {
    "wspd": u"m/s",
    "wspd_max": u"m/s",
    "wdir": u"°",
    "wpd": u"W/m2",
    "t": u"°C",
    "t_max": u"°C",
    "t_min": u"°C",
    "psfc": u"hPa",
    "slp": u"hPa",
    "qv": u"g/kg",
    "summary": u"",
    "rose": u"‰",
    "wpdrose": u"kWh/m²",
    "dist": u"%",
}

# 高度的字典
merra2_levels_d = {
    "wspd": ["50", "10", "2"],
    "wspd_max": ["50", "10", "2"],
    "wdir": ["50", "10", "2"],
    "wpd": ["50"],
    "t": ["10", "2", "sfc"],
    "t_max": ["10", "2", "sfc"],
    "t_min": ["10", "2", "sfc"],
    "psfc": ["sfc"],
    "slp": ["sfc"],
    "qv": ["10", "2"],
    "summary": ['all'],
    "rose": ["50", "10", "2"],
    "wpdrose": ["50"],
    "dist": ["50", "10", "2"],
}

#默认高度的字典
merra2_levels_default = {
    "wspd": ["50"],
    "wspd_max": ["50"],
    "wdir": ["50"],
    "wpd": ["50"],
    "t": ["2"],
    "t_max": ["2"],
    "t_min": ["2"],
    "psfc": ["sfc"],
    "slp": ["sfc"],
    "qv": ["2"],
    "summary": ['mixed'],
    "rose": ["50"],
    "wpdrose": ["50"],
    "dist": ["50"],
}


# 推荐的色表
merra2_suggest_cmap_d = {
    't': 'temperature',
    't_max': 'temperature',
    't_min': 'temperature',
    'psfc': 'pressure',
    'slp': 'pressure',
    'qv': 'rain'
}


merra2_suggest_range_d = {
    "wspd": (3.0, 10.0),
    "wspd_max": (10.0, 35.0),
    "wdir": (0.0, 360.0),
    "wpd": (0.0, 1000.0),
    "t": (-10.0, 30.0),
    "t_max": (5.0, 45.0),
    "t_min": (-40.0, 20.0),
    "psfc": (None, None),
    "slp": (1000, 1020),
    "qv": (0.0, 22.0)
}


def latlon_to_merra2_jyix(lat, lon):
    jy = np.round((lat - (-90)) / merra2_dy)
    ix = np.round((lon - (-180)) / merra2_dx)
    return '%03d_%03d' % (jy, ix)


def get_merra2_rect(lon1, lat1, lon2, lat2):
    latmin = min(lat1, lat2); latmax = max(lat1, lat2)
    lonmin = min(lon1, lon2); lonmax = max(lon1, lon2)
    latmin -= merra2_dy; latmax += merra2_dy
    lonmin -= merra2_dx; lonmax += merra2_dx
    jymin, ixmin = latlon_to_merra2_jyix(latmin, lonmin).split('_')
    jymax, ixmax = latlon_to_merra2_jyix(latmax, lonmax).split('_')

    res = {
        'jy_ix': '%s:%s_%s:%s' % (jymin, jymax, ixmin, ixmax)
    }
    return res


def parse_merra2_uri(uri):
    result = {}
    m = re.match(r'(?P<dataset>[^/]+)/(?P<subset>[^/]+)/(?P<varname>[^/]+)/(?P<time>[^-+Z/]+)(Z?)(?P<timezone>[-+0-9]*)/(?P<level>[^/]+)/(?P<jy_ix>[^/]+)$', uri)
    if not m:
        raise ValueError('%s is not a proper merra2 uri' % uri)
    for field in ['dataset', 'subset', 'varname', 'time', 'level', 'jy_ix']:
        result[field] = m.group(field)

    time_str = m.group('time')
    timezone_str = m.group('timezone')
    try:
        timezone = int(timezone_str)
    except ValueError:
        timezone = 0

    if ',' in time_str:
        times = time_str.split(',')
        result['times'] = times
        result['timetype'] = 'year_list'
    elif ':' in time_str:
        begdt, enddt, tdelta = parse_slice_str(time_str, default_step='1h')
        begdt = T(begdt)
        enddt = T(enddt)
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = tdelta
        result['timetype'] = 'dtrange'
    elif len(m.group('time')) == 4:  # year
        begdt = T(m.group('time')+'0101')
        enddt = begdt + TD('1Y')
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = '1h'
        result['timetype'] = 'year'
    else:
        raise ValueError('%s is not a proper merra2 uri' % uri)
    result['timezone'] = timezone

    return result


def merra2_yearlymean2summary(orig):
    tags = orig['tags']
    sub_varnames = ['ws50m', 'wpd', 't10m', 't2m', 'ts', 'ps', 'slp']
    info = {
        'sub_varnames': sub_varnames,
        'type': 'merra_summary',
    }
    data = {}
    for vn in sub_varnames:
        rawvar = orig['contents'][vn]
        subvar = {
            'name': vn,
            'zh_name': merra2_zh_varname_d.get(vn, vn),
            'values': rawvar['data']['values'],
            'std': rawvar['data']['std'],
            'count': rawvar['data']['count'],
            'units': rawvar['data']['units'],
        }
        data[vn] = subvar
    common = orig['common'].copy()
    common['coords'] = {}
    for key in ['dataset', 'subset', 'time', 'jy_ix']:
        common['coords'][key] = common['info'].pop(key, '')

    res = {
        'uri': orig['uri'],
        'common': common,
        'tags': {
            'data': orig['uri']
        },
        'contents': {
            'data': {
                'info': info,
                'data': data,
                'coords': {
                    'varname': 'mean',
                }
            },
        }
    }

    return res


def merra2_py_add_kwargs(get_data_func):
    def modified_func(instance, uri, request=None, **kwargs):
        try:
            user_profile = request.user.profile
        except Exception:
            user_profile = None

        uri_info = instance.parse_uri(uri)
        varname = uri_info['varname']
        key = "dataset:{}:summary_sub_vars".format(instance.topset.name)
        if user_profile:
            sub_varnames = user_profile.get_setting(key, instance.py_summary_sub_varnames)
        else:
            sub_varnames = instance.py_summary_sub_varnames
        kwargs['sub_varnames'] = ','.join(sub_varnames)
        return get_data_func(instance, uri, request=request, **kwargs)
    return modified_func


def merra2_point_grid_geometry(get_data_func):
    """给Merra2 Point系列 get_data接口加的装饰器, 用来给结果的info属性添加grid_geometry信息"""
    def modified_func(instance, uri, request=None, **kwargs):
        res = get_data_func(instance, uri, request=request, **kwargs)
        if 'uri_info' not in kwargs or kwargs['uri_info'] is None:
            kwargs['uri_info'] = parse_merra2_uri(uri)
        uri_info = kwargs['uri_info']
        if isinstance(res, dict) and res.get('type') in ("dataunit", "datapack"):
            jy, ix = [int(tk) for tk in uri_info['jy_ix'].split('_')]
            pts = []
            for dj, di in ((-0.5, -0.5), (-0.5, 0.5), (0.5, 0.5), (0.5, -0.5)):
                pts.append(((ix+di)*merra2_dx-180.0, (jy+dj)*merra2_dy-90.0))
            pts.append(pts[0])
            if res.get('type') == "dataunit":
                dataunits = [res]
            elif res.get('type') == "datapack":
                dataunits = res['contents'].values()

            for dataunit in dataunits:
                dataunit["info"]["grid_geometry"] = {
                    "type": "Polygon",
                    "coordinates": [pts]
                }
        return res
    return modified_func
